import os
import torch
import torch.nn.functional as F
import tkinter as tk
from PIL import Image, ImageDraw
from torchvision import transforms


# 导入自定义模块
from dataloader import get_paths
from model import LeNet5

# 获取数据和输出路径
_, output_dir = get_paths()

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class DrawingApp:
    """手写数字识别绘图界面"""
    
    def __init__(self, model):
        self.model = model
        self.model.eval()  # 设置为评估模式
        
        self.root = tk.Tk()
        self.root.title("手写数字识别 - LeNet-5")
        
        # 设置窗口大小
        self.root.geometry("600x450")
        
        # 创建画布框架
        self.canvas_frame = tk.Frame(self.root, bd=2, relief=tk.SUNKEN)
        self.canvas_frame.grid(row=0, column=0, padx=10, pady=10)
        
        # 创建画布
        self.canvas = tk.Canvas(self.canvas_frame, width=280, height=280, bg="black")
        self.canvas.pack()
        
        # 绑定鼠标事件
        self.canvas.bind("<B1-Motion>", self.paint)
        
        # 创建控制按钮框架
        self.control_frame = tk.Frame(self.root)
        self.control_frame.grid(row=0, column=1, padx=10, pady=10)
        
        # 识别结果标签
        self.result_label = tk.Label(self.control_frame, text="绘制一个数字", font=("Helvetica", 16))
        self.result_label.pack(pady=10)
        
        # 预测结果
        self.prediction_text = tk.StringVar()
        self.prediction_text.set("预测结果: ")
        self.prediction_label = tk.Label(self.control_frame, textvariable=self.prediction_text, font=("Helvetica", 14))
        self.prediction_label.pack(pady=5)
        
        # 概率条形图框架
        self.probs_frame = tk.Frame(self.control_frame)
        self.probs_frame.pack(pady=10, fill=tk.X)
        
        # 概率标签
        self.prob_labels = []
        for i in range(10):
            frame = tk.Frame(self.probs_frame)
            frame.pack(fill=tk.X, pady=2)
            
            # 数字标签
            digit_label = tk.Label(frame, text=f"{i}:", width=2)
            digit_label.pack(side=tk.LEFT, padx=5)
            
            # 概率条
            prob_bar = tk.Canvas(frame, width=150, height=15, bg="white", bd=1, relief=tk.SUNKEN)
            prob_bar.pack(side=tk.LEFT)
            
            # 概率值标签
            prob_value = tk.Label(frame, text="0%", width=5)
            prob_value.pack(side=tk.LEFT, padx=5)
            
            self.prob_labels.append((prob_bar, prob_value))
        
        # 清除按钮
        self.clear_button = tk.Button(self.control_frame, text="清除", command=self.clear_canvas)
        self.clear_button.pack(pady=10)
        
        # 识别按钮
        self.recognize_button = tk.Button(self.control_frame, text="识别", command=self.recognize)
        self.recognize_button.pack(pady=5)
        
        # 退出按钮
        self.quit_button = tk.Button(self.control_frame, text="退出", command=self.root.quit)
        self.quit_button.pack(pady=5)
        
        # 画笔宽度
        self.line_width = 15
        
        # 上一个绘制点
        self.x = None
        self.y = None
        
        # 保存用于识别的图像
        self.image = Image.new("L", (280, 280), "black")
        self.draw = ImageDraw.Draw(self.image)
    
    def paint(self, event):
        """处理鼠标绘制事件"""
        # 绘制线条
        if self.x is not None and self.y is not None:
            self.canvas.create_line(self.x, self.y, event.x, event.y,
                                   fill="white", width=self.line_width, capstyle=tk.ROUND, smooth=tk.TRUE)
            self.draw.line([self.x, self.y, event.x, event.y], fill="white", width=self.line_width)
        
        self.x = event.x
        self.y = event.y
    
    def clear_canvas(self):
        """清除画布"""
        self.canvas.delete("all")
        self.image = Image.new("L", (280, 280), "black")
        self.draw = ImageDraw.Draw(self.image)
        self.x = None
        self.y = None
        
        # 重置预测结果
        self.prediction_text.set("预测结果: ")
        
        # 重置概率条
        for prob_bar, prob_value in self.prob_labels:
            prob_bar.delete("all")
            prob_value.config(text="0%")
    
    def recognize(self):
        """识别绘制的数字"""
        # 将图像调整为模型输入大小
        img = self.image.resize((28, 28))
        
        # 保存绘制的图像用于调试
        debug_path = os.path.join(output_dir, 'last_drawing.png')
        img.save(debug_path)
        
        # 转换为模型所需的格式
        img_tensor = transforms.ToTensor()(img)
        img_tensor = transforms.Normalize((0.1307,), (0.3081,))(img_tensor)
        img_tensor = img_tensor.unsqueeze(0).to(device)
        
        # 进行预测
        with torch.no_grad():
            output = self.model(img_tensor)
            probabilities = F.softmax(output, dim=1)[0]
            predicted_class = torch.argmax(output, dim=1).item()
        
        # 显示预测结果
        self.prediction_text.set(f"预测结果: {predicted_class}")
        
        # 更新概率条
        for i, (prob_bar, prob_value) in enumerate(self.prob_labels):
            prob = probabilities[i].item()
            width = int(prob * 150)
            
            # 清除旧条形
            prob_bar.delete("all")
            
            # 绘制新条形
            if i == predicted_class:
                color = "green"
            else:
                color = "blue"
            
            prob_bar.create_rectangle(0, 0, width, 15, fill=color, outline="")
            prob_value.config(text=f"{prob:.1%}")
    
    def run(self):
        """运行GUI应用"""
        self.root.mainloop()

def main():
    """主函数"""
    # 加载模型
    model_path = os.path.join(output_dir, 'lenet5_model.pth')
    
    if not os.path.exists(model_path):
        print(f"未找到已训练的模型 {model_path}，请先运行train.py训练模型")
        return
    
    # 创建模型实例
    model = LeNet5().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    
    print("启动手写数字识别界面...")
    app = DrawingApp(model)
    app.run()

if __name__ == "__main__":
    main()